import sklearn
from sklearn.datasets import load_iris, fetch_20newsgroups, load_files
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import GaussianNB
from sklearn.metrics import accuracy_score
import os
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB


# 获取安装目录
def default_path():
    # 获取 Scikit-learn 安装目录
    sklearn_dir = os.path.dirname(sklearn.__file__)
    # 拼接数据集路径
    # data_path = os.path.join(sklearn_dir, 'datasets', 'data', 'iris.csv')
    print("数据集绝对路径:", sklearn_dir)


# 高斯贝叶斯算法
def gaussian_nb():
    # 加载数据集
    data = load_iris()
    X, y = data.data, data.target
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    # 初始化高斯朴素贝叶斯
    model = GaussianNB()
    # 训练模型
    model.fit(X_train, y_train)
    # 预测
    y_pred = model.predict(X_test)
    # 评估准确率
    print("Accuracy:", accuracy_score(y_test, y_pred))


# 多项式朴素贝叶斯（文本分类）
def multinomial_bayes():
    # 示例文本数据（邮件分类：0=正常邮件，1=垃圾邮件）
    texts = ["free money now", "hello world", "buy free prize", "hello friend"]
    labels = [1, 0, 1, 0]

    # 特征工程——将文本转换为词频向量
    vectorizer = CountVectorizer()
    X = vectorizer.fit_transform(texts)

    # 训练模型
    model = MultinomialNB()
    model.fit(X, labels)

    # 预测新文本
    new_text = ["free prize"]
    new_vector = vectorizer.transform(new_text)
    print("Prediction:", model.predict(new_vector))  # 输出: [1]（垃圾邮件）


def nb_news_text_classify():
    """
    用朴素贝叶斯算法对新闻进行分类
    步骤：
    1. 加载数据
    2. 数据预处理（TF-IDF向量化）
    3. 构建模型（多项式朴素贝叶斯）
    4. 训练和评估
    5. 可视化结果
    """
    # 1.获取数据
    news = fetch_20newsgroups(subset='all')
    # 查看类别和样本数
    print("类别和样本数：\n", news.target_names, news.target.shape)

    # 2. 数据拆分 -------------------------------------------------
    X_train, X_test, y_train, y_test = train_test_split(news.data, news.target)
    # print(f"训练集：{X_train}")

    # 3. 特征工程
    transfer = TfidfVectorizer()
    X_train = transfer.fit_transform(X_train)
    X_test = transfer.transform(X_test)

    # 4. 朴素贝叶斯算法预估器
    estimator = MultinomialNB()  # 多项式朴素贝叶斯（文本分类）
    estimator.fit(X_train, y_train)

    # 5. 模型评估
    y_predict = estimator.predict(X_test)
    print("预测结果为：", y_predict)
    print(f"直接比对真实值和预测值：\n{y_test == y_predict}")
    print(f"准确率：\n{estimator.score(X_test, y_test)}")


if __name__ == '__main__':
    # NB 是 Naive Bayes（朴素贝叶斯） 的缩写
    # gaussian_nb()
    # multinomial_bayes()
    nb_news_text_classify()
